#!/usr/bin/env python3
"""
GSI-EnKF Comparison Tool
Compares GSI and EnKF stdout files to identify integration issues
"""

import sys
import json
import os
import copy
from typing import Dict, List, Tuple
from .analyze_gsi_stdout import GSIAnalyzer
from .analyze_enkf_stdout import EnKFAnalyzer


class GSIEnKFComparator:
    """Compare GSI and EnKF runs to identify integration issues"""
    
    def __init__(self, gsi_file: str, enkf_file: str):
        self.gsi_analyzer = GSIAnalyzer(gsi_file)
        self.enkf_analyzer = EnKFAnalyzer(enkf_file)
        self.comparison = {
            'compatibility_issues': [],
            'configuration_mismatches': [],
            'data_flow_problems': [],
            'recommendations': [],
            'ensemble_status': {},
            'observation_comparison': {}
        }
        
    def analyze_both(self):
        """Run analysis on both files"""
        # Analyze GSI
        self.gsi_analyzer.read_file()
        self.gsi_analyzer.analyze()
        
        # Analyze EnKF
        self.enkf_analyzer.read_file()
        self.enkf_analyzer.analyze()
        
    def compare_execution_status(self):
        """Compare execution status of both systems"""
        gsi_status = self.gsi_analyzer.analysis['execution_status']
        enkf_status = self.enkf_analyzer.analysis['execution_status']
        
        gsi_complete = gsi_status.get('completed', False)
        enkf_complete = enkf_status.get('completed', False)
        
        if gsi_complete and not enkf_complete:
            self.comparison['compatibility_issues'].append(
                "GSI completed successfully but EnKF failed - hybrid system broken"
            )
        elif not gsi_complete and enkf_complete:
            self.comparison['compatibility_issues'].append(
                "EnKF completed but GSI failed - unusual configuration"
            )
            
    def check_ensemble_integration(self):
        """Check if GSI is using EnKF ensemble"""
        gsi_ensemble = self.gsi_analyzer.analysis.get('ensemble_info', {})
        enkf_config = self.enkf_analyzer.analysis['configuration']
        enkf_ensemble = self.enkf_analyzer.analysis['ensemble_info']
        
        # Get ensemble size from EnKF
        enkf_size = enkf_config.get('ensemble_size', 0)
        
        # Store ensemble status
        self.comparison['ensemble_status'] = {
            'enkf_size': enkf_size,
            'gsi_uses_ensemble': gsi_ensemble.get('uses_ensemble', False),
            'gsi_ensemble_members': gsi_ensemble.get('n_members', 0),
            'enkf_has_spread': not enkf_ensemble.get('identical_members', False)
        }
        
        # Check for configuration issues
        if enkf_size > 1:
            if not gsi_ensemble.get('uses_ensemble'):
                self.comparison['configuration_mismatches'].append(
                    f"EnKF has {enkf_size} members but GSI not using ensemble"
                )
                self.comparison['recommendations'].append(
                    "Configure GSI for hybrid mode with ensemble background"
                )
                
            # Check member count consistency
            gsi_members = gsi_ensemble.get('n_members', 0)
            if gsi_members > 0 and gsi_members != enkf_size:
                self.comparison['configuration_mismatches'].append(
                    f"Ensemble size mismatch: GSI expects {gsi_members} members but EnKF has {enkf_size}"
                )
                
        # Check for zero spread issue
        if enkf_ensemble.get('identical_members'):
            self.comparison['data_flow_problems'].append(
                "EnKF has zero ensemble spread - all members identical"
            )
            self.comparison['recommendations'].append(
                "Provide unique ensemble member backgrounds to EnKF"
            )
            
    def check_observation_consistency(self):
        """Check observation usage between systems"""
        gsi_obs = self.gsi_analyzer.analysis['observations']
        enkf_obs = self.enkf_analyzer.analysis['observations']
        
        # Store comparison data
        self.comparison['observation_comparison'] = {
            'gsi_total': gsi_obs.get('total_used', 0),
            'gsi_types': len(gsi_obs.get('types', {})),
            'enkf_total': enkf_obs.get('total', 0),
            'enkf_used': enkf_obs.get('used', 0),
            'enkf_types': len(enkf_obs.get('types_configured', {}))
        }
        
        # Compare observation counts
        if gsi_obs.get('total_used', 0) > 0 and enkf_obs.get('used', 0) == 0:
            self.comparison['data_flow_problems'].append(
                f"GSI used {gsi_obs['total_used']:,} observations but EnKF used 0"
            )
            
        # Check for significant differences in observation counts
        enkf_total = enkf_obs.get('total', 0)
        gsi_total = gsi_obs.get('total_used', 0)
        if enkf_total > 0 and abs(gsi_total - enkf_total) > 100:
            self.comparison['data_flow_problems'].append(
                f"Observation count mismatch: GSI used {gsi_total:,} vs EnKF processed {enkf_total:,}"
            )
            
        # Check radiance usage
        gsi_rad = self.gsi_analyzer.analysis.get('radiance_info', {})
        if not gsi_rad.get('has_radiances'):
            self.comparison['configuration_mismatches'].append(
                "No radiance observations in GSI - limited to conventional obs"
            )
            if gsi_obs.get('total_used', 0) < 10000:
                self.comparison['recommendations'].append(
                    "Configure satellite radiance assimilation to increase observation coverage"
                )
                
    def check_grid_consistency(self):
        """Check grid configuration consistency"""
        gsi_grid = self.gsi_analyzer.analysis['grid_info']
        enkf_grid = self.enkf_analyzer.analysis['grid_info']
        
        # Check dimensions
        for dim in ['nx', 'ny', 'nz']:
            gsi_val = gsi_grid.get(dim)
            enkf_val = enkf_grid.get(dim)
            if gsi_val and enkf_val and gsi_val != enkf_val:
                self.comparison['configuration_mismatches'].append(
                    f"Grid {dim} mismatch: GSI has {gsi_val} but EnKF has {enkf_val}"
                )
                
    def check_time_consistency(self):
        """Check if both systems are analyzing the same time"""
        gsi_time = self.gsi_analyzer.analysis['execution_status'].get('analysis_time')
        enkf_time = self.enkf_analyzer.analysis['configuration'].get('analysis_time')
        
        if gsi_time and enkf_time and gsi_time != enkf_time:
            self.comparison['compatibility_issues'].append(
                f"Time mismatch: GSI analyzing {gsi_time}, EnKF analyzing {enkf_time}"
            )
            
    def check_analysis_quality(self):
        """Check analysis increment quality"""
        gsi_inc = self.gsi_analyzer.analysis.get('increments', {})
        enkf_inc = self.enkf_analyzer.analysis.get('increments', {})
        
        # Check if increments are too small (possible issue)
        if gsi_inc:
            small_increments = []
            for var, stats in gsi_inc.items():
                if abs(stats.get('mean', 0)) < 1e-10:
                    small_increments.append(var)
            
            if len(small_increments) > 3:
                self.comparison['data_flow_problems'].append(
                    "GSI increments near zero - possible data assimilation issue"
                )
                
        # Check EnKF increments
        if enkf_inc:
            enkf_small = []
            for var, stats in enkf_inc.items():
                if 'max' in stats and abs(stats['max']) < 1e-7:
                    enkf_small.append(var)
                    
            if len(enkf_small) > 2:
                self.comparison['data_flow_problems'].append(
                    "EnKF increments near zero - no effective update performed"
                )
                
    def generate_comparison_report(self) -> str:
        """Generate comparison report"""
        report = []
        report.append("=" * 80)
        report.append("GSI-EnKF INTEGRATION ANALYSIS")
        report.append("=" * 80)
        report.append("")
        
        # Summary
        report.append("SYSTEM STATUS:")
        gsi_complete = self.gsi_analyzer.analysis['execution_status'].get('completed', False)
        enkf_complete = self.enkf_analyzer.analysis['execution_status'].get('completed', False)
        
        report.append(f"  GSI:  {'COMPLETED' if gsi_complete else 'FAILED/INCOMPLETE'}")
        report.append(f"  EnKF: {'COMPLETED' if enkf_complete else 'FAILED/INCOMPLETE'}")
        report.append("")
        
        # Compatibility Issues
        if self.comparison['compatibility_issues']:
            report.append("COMPATIBILITY ISSUES:")
            for issue in self.comparison['compatibility_issues']:
                report.append(f"  ⚠️  {issue}")
            report.append("")
            
        # Configuration Mismatches
        if self.comparison['configuration_mismatches']:
            report.append("CONFIGURATION MISMATCHES:")
            for issue in self.comparison['configuration_mismatches']:
                report.append(f"  ❌ {issue}")
            report.append("")
            
        # Data Flow Problems
        if self.comparison['data_flow_problems']:
            report.append("DATA FLOW PROBLEMS:")
            for issue in self.comparison['data_flow_problems']:
                report.append(f"  🔄 {issue}")
            report.append("")
            
        # Key Metrics Comparison
        report.append("KEY METRICS:")
        
        # GSI metrics
        gsi_config = self.gsi_analyzer.analysis['configuration']
        gsi_obs = self.gsi_analyzer.analysis['observations']
        gsi_ensemble = self.gsi_analyzer.analysis.get('ensemble_info', {})
        
        report.append("  GSI:")
        report.append(f"    Observations used: {gsi_obs.get('total_used', 0):,}")
        report.append(f"    Observation types: {len(gsi_obs.get('types', {}))}")
        report.append(f"    Variational mode: {gsi_config.get('variational_mode', 'Unknown')}")
        report.append(f"    Using ensemble: {'Yes' if gsi_ensemble.get('uses_ensemble') else 'No'}")
        
        # EnKF metrics
        enkf_config = self.enkf_analyzer.analysis['configuration']
        enkf_obs = self.enkf_analyzer.analysis['observations']
        enkf_ens = self.enkf_analyzer.analysis['ensemble_info']
        
        report.append("  EnKF:")
        report.append(f"    Observations used: {enkf_obs.get('used', 0):,}")
        report.append(f"    Ensemble size: {enkf_config.get('ensemble_size', 'unknown')}")
        report.append(f"    Has ensemble spread: {'No' if enkf_ens.get('identical_members') else 'Yes'}")
        report.append("")
        
        # Ensemble Integration
        ens_status = self.comparison['ensemble_status']
        if ens_status:
            report.append("ENSEMBLE INTEGRATION:")
            report.append(f"  EnKF ensemble size: {ens_status['enkf_size']}")
            report.append(f"  GSI using ensemble: {'Yes' if ens_status['gsi_uses_ensemble'] else 'No'}")
            report.append(f"  EnKF has spread: {'Yes' if ens_status['enkf_has_spread'] else 'No'}")
            report.append("")
        
        # Recommendations
        if self.comparison['recommendations']:
            report.append("RECOMMENDATIONS:")
            for i, rec in enumerate(self.comparison['recommendations'], 1):
                report.append(f"  {i}. {rec}")
            report.append("")
            
        # Overall Assessment
        report.append("OVERALL ASSESSMENT:")
        if gsi_complete and enkf_complete and not any([
            self.comparison['compatibility_issues'],
            self.comparison['configuration_mismatches'],
            self.comparison['data_flow_problems']
        ]):
            report.append("  ✅ Both systems completed successfully with no major issues detected")
        else:
            n_issues = (len(self.comparison['compatibility_issues']) + 
                       len(self.comparison['configuration_mismatches']) + 
                       len(self.comparison['data_flow_problems']))
            report.append(f"  ❌ Hybrid EnKF-GSI system has {n_issues} integration problems requiring attention")
            
        report.append("")
        report.append("=" * 80)
        
        return '\n'.join(report)
        
    def run(self):
        """Run the complete comparison"""
        self.analyze_both()
        self.compare_execution_status()
        self.check_ensemble_integration()
        self.check_observation_consistency()
        self.check_grid_consistency()
        self.check_time_consistency()
        self.check_analysis_quality()
        return self.generate_comparison_report()


def main():
    """Main function"""
    args = sys.argv[1:]

    # Process optional arguments
    json_file = None
    if '--json' in args:
        try:
            json_idx = args.index('--json')
            json_file = args.pop(json_idx + 1)
            args.pop(json_idx)
        except IndexError:
            print("Usage: python compare_gsi_enkf.py [<gsi_stdout> <enkf_stdout>] [--json output.json]")
            sys.exit(1)

    # Process positional arguments
    if len(args) == 2:
        gsi_file, enkf_file = args
    elif len(args) == 0:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        gsi_file = os.path.join(script_dir, 'gsi_stdout')
        enkf_file = os.path.join(script_dir, 'enkf_stdout')
    else:
        print("Usage: python compare_gsi_enkf.py [<gsi_stdout> <enkf_stdout>] [--json output.json]")
        sys.exit(1)

    try:
        comparator = GSIEnKFComparator(gsi_file, enkf_file)
        report = comparator.run()
        print(report)

        if json_file:
            output_data = {
                'gsi_analysis': copy.deepcopy(comparator.gsi_analyzer.analysis),
                'enkf_analysis': copy.deepcopy(comparator.enkf_analyzer.analysis),
                'comparison': comparator.comparison
            }
            # Filter out warnings from the 'issues' list for GSI analysis
            if 'issues' in output_data['gsi_analysis']:
                output_data['gsi_analysis']['issues'] = [
                    issue for issue in output_data['gsi_analysis']['issues']
                    if issue.get('type') != 'warning'
                ]
            # Filter out warnings from the 'issues' list for EnKF analysis
            if 'issues' in output_data['enkf_analysis']:
                output_data['enkf_analysis']['issues'] = [
                    issue for issue in output_data['enkf_analysis']['issues']
                    if issue.get('type') != 'warning'
                ]

            with open(json_file, 'w') as f:
                json.dump(output_data, f, indent=2)
            print(f"\nJSON comparison saved to: {json_file}")

    except Exception as e:
        print(f"Error during comparison: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()